import pdfplumber
import fitz
import re
import logging 
import os


path = r"pdf_files/000663_永安木业2023年年度报告.pdf" #初始pdf

#截取 节 财务报告后面的页数
def getmulu():
    pdff = fitz.open(path)
    total_pages = pdff.page_count
    toc = pdff.get_toc()
    for text in toc:
        match=re.findall(r"节\s+财务报告",text[1])
        if len(match) > 0:
            start_page = max(0, text[2] - 1)
            new_doc = fitz.open()
            for page_num in range(start_page, total_pages):
                new_doc.insert_pdf(pdff, from_page=page_num, to_page=page_num)
            path2=f'财务报告{text[2]}_{total_pages}.pdf'
            # new_doc.save(f'财务报告{text[2]}_{total_pages}.pdf')
            if os.path.exists(path2):
                new_doc.save(f'财务报告{text[2]}_{total_pages}.pdf')           
            new_doc.close()
    pdff.close()
    return path2


def checkWord(path2,identifyWord):
    pdff = pdfplumber.open(path2)
    list =[]
    returnList=[]
    for page in pdff.pages:
        for words in page.extract_words():
            if identifyWord in words.get('text'):
                list.append((page.page_number-1,words.get('top'),words.get('text')))
    if len(list)==0:
        logging.info('没有找到指定的关键词')
    elif len(list)>1:
        i=1
        for word in list:
            if len(word[2]) <20:
                table =getTableByPage(pdff,word[0])
                if len(table) > 0:
                    # print(f'----找到的关键有多个，第{i}个在{word[0]+1}页,全文为:{word[2]}')
                    # print(f'----表格为{table}')
                    returnList.append(table)
                    i+=1          
    pdff.close()
    return returnList

def getTableByPage(pdff,pageNum):
    # pdff = pdfplumber.open(path)
    page = pdff.pages[pageNum]
    table = page.extract_tables()
    tablelist=[]
    if len(table)>0:
        tablelist.append(table[-1])
        for i in range(10):
            pageNum+=1
            page = pdff.pages[pageNum]
            table = page.extract_tables()
            if len(table) == 1:
                tablelist.append(table[0])
            elif len(table) > 1:
                tablelist.append(table[0])
                break
            else:
                break
    
    return tablelist

#清洗数据，去掉多层和数据中的换行符
def qinxidarta(list):
    wordList=[]
    for everyList in list[0]:
        for everyList1 in everyList:
            for n in range(len(everyList1)):
                if everyList1[n] != None:
                    everyList1[n]=everyList1[n].replace("\n", "")
            wordList.append(everyList1)
    return wordList   

def merge_code(keyword):
    path2=getmulu()#截取
    # list=checkWord(path2,'公司资产负债表')#要查找关键词
    list=checkWord(path2,keyword)
    list1=qinxidarta(list)#清洗数据，去掉多层和数据中的换行符
    return list1


if __name__=='__main__':
    path2=getmulu()#截取
    # list=checkWord(path2,'公司资产负债表')#要查找关键词
    list=checkWord(path2,'合并利润表')
    list1=qinxidarta(list)#清洗数据，去掉多层和数据中的换行符
    print(list1)
    
    